package com.example.aicode.chat.rag;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import jakarta.annotation.Resource;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.List;

/**
 * RAG 配置类
 */
@Configuration
public class RagConfig {

    @Resource
    private EmbeddingModel embeddingModel;

    @Resource
    private EmbeddingStore<TextSegment> embeddingStore;

    @Bean
    public ContentRetriever contentRetriever() {
        // 1 加载文档
        List<Document> documents = FileSystemDocumentLoader
                .loadDocuments("src/main/resources/markdown");
        // 2 文档切割，最大1000个字符，每次最多重叠200个
        DocumentByParagraphSplitter splitter = new DocumentByParagraphSplitter(1000, 200);
        // 3 自定义文档加载器，将文档切割成向量并保存到数据库中
        EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
                .documentSplitter(splitter)
                .textSegmentTransformer(textSegment -> TextSegment.from(
                        textSegment.metadata().getString("fileName") + "\n" + textSegment.text(),
                        textSegment.metadata()))
                .embeddingModel(embeddingModel).embeddingStore(embeddingStore).build();
        // 4 加载文档
        ingestor.ingest(documents);
        // 5 自定义内容加载器
        EmbeddingStoreContentRetriever embeddingStoreContentRetriever = EmbeddingStoreContentRetriever
                .builder().embeddingStore(embeddingStore).embeddingModel(embeddingModel)
                .maxResults(5) // 最多返回的结果数
                .minScore(0.75) // 过滤分数小于0.75的结果
                .build();
        return embeddingStoreContentRetriever;
    }
}
